#!/usr/bin/python

"""
ssh-fwknop (c) Rich Lundeen

Licensed under GPLv3

This is spaghetti code by necessity. 
The problem is simple enough, to separate the options for fwknop and ssh, and this code
accomplishes it. Hopefully in a fairly maintainable way. However, there are a lot of kludges...

Still, it's usefuly to people like me who frequently forget options or who don't want to type in 
a hostname twice.

"""

import getopt, sys, ConfigParser, os, subprocess
import ssh_fwopts

class getoptions:
  #longopts_args, longopt_noargs, shortopt_args, shortop_noargs are passed exactly as if to getopt
  def __init__(self, config=None, section=None, commandlineoptions="", longopt_args=[], longopt_noargs=[],
               shortopt_args="", shortopt_noargs="", host="", binpath=""):
    self.finoptions = []
    self.finargs    = []
    #host is gotten from 1) -D options in fwknop 2) argv[0] 3) host in config file
    self.config=config
    self.host = host
    self.binpath = binpath
    self.section=section
 
    #these are special arguments that need to be parsed by ssh
    #the second value indicates if it needs to be searched for or not by listify
    self.fwknop_opts = ([], False)
    self.__separgs__(commandlineoptions, longopt_args, longopt_noargs, shortopt_args, shortopt_noargs)
    self.__configread__(section, longopt_args, longopt_noargs, shortopt_args, shortopt_noargs)

  def usage(self):
    sys.stderr.write("Usage: ssh-fwknop [--fwknop='fwknop options'] [--ssh | --sftp | --scp] \n " +
                     "                  [sh-options] [user@]hostname [command]\n")
    sys.exit(-2)

  """This function is necessary because ssh doesn't care where you place the arguments, whereas
     getopt by default does. eg with getopt ssh myserver.com -p 33 would treat myserver.com -p 33
     as the arguments, while ssh still treats -p 33 as options """
  def __separgs__(self, commandlineoptions, longopt_args, longopt_noargs, shortopt_args, shortopt_noargs):
    args_toparse = True
    args = commandlineoptions
    while args_toparse:
      #because ssh can have only certain arguments, and options can be separated,
      try:
        opts, args = getopt.getopt(args, (shortopt_args + shortopt_noargs), (longopt_args + longopt_noargs))
      except getopt.GetoptError, err:
        sys.stderr.write( str(err))
        self.usage()
      self.finoptions = self.finoptions[:] + opts
      #see if there's anything left to parse to grap the real arguments
      args_toparse = False
      morereal = False
      for i in args:
        if "-" in i:
          args_toparse = True
          self.finargs.append(args[0])
          args = args[1:]
          morereal = True
          break
      if not morereal:
        self.finargs = self.finargs[:] + args[:]

  def __configread__(self, section, longopt_args, longopt_noargs, shortopt_args, shortopt_noargs):
    #get options from the configuration file
    for i in longopt_args:
      try:
        thisoption = self.config.get(section, i) 
        i = "--"+i
        #command line has precidance
        if not self.__isinoptions__(i):
          self.finoptions.append((i, thisoption))
      except ConfigParser.NoSectionError:
        #no section in the config file, so need for the loop
        break
      except (NameError, ConfigParser.NoOptionError):
        continue
      except AttributeError:
        continue
    for i in longopt_noargs:
      try:
        thisoption = self.config.get(section, i)
        i = "--"+i
        if not self.__isinoptions__(i) and strip(thisoption).lower == "true":
          self.finoptions.append((i,''))
      except ConfigParser.NoSectionError:
        #no section in the config file, so need for the loop
        break
      except (NameError, ConfigParser.NoOptionError):
        continue
      except AttributeError:
        continue
    for i in shortopt_args:
      try:
        thisoption = self.config.get(section, i) 
        i = "-" + i
        if not self.__isinoptions__(i):
        #command line has precidance
          self.finoptions.append((i, thisoption))
      except ConfigParser.NoSectionError:
        #no section in the config file, so need for the loop
        break
      except (NameError, ConfigParser.NoOptionError):
        continue
      except AttributeError:
        continue
    for i in shortopt_noargs:
      try:
        thisoption = self.config.get(section, i)
        i = "-" + i
        if (not self.__isinoptions__(i)) and (thisoption.strip().lower() == "true"):
          self.finoptions.append((i,''))
      except ConfigParser.NoSectionError:
        #no section in the config file, so need for the loop
        break
      except (NameError, ConfigParser.NoOptionError):
        continue
      except AttributeError:
        continue
    
  def __isinoptions__(self, option):
    for i in self.finoptions:
      if i[0].strip() == option.strip():
        return True
    return False 

  def listify(self):
    totalstring = [self.binpath]
    for i,j in self.finoptions:
      if "fwknop" in i:
        self.fwknop_opts = (j.split(), True)
        continue
      totalstring.append(i) 
      if j != "":
        totalstring.append(j) 
    #TODO make this more generic
    if (self.section == "ssh" or self.section == "sftp" or self.section == "scp")and self.finargs == []:
      self.finargs.append(self.gethost())
    return totalstring + self.finargs

  def get_fwknop(self):
    #listify should actually just be called first
    if self.fwknop_opts[1] == False:
      self.listify()
    return self.fwknop_opts[0]

  #gethost only ever should return the host
  #TODO important, handle special characters like @,: and scp which is different 
  def gethost(self):
    if self.host == "":
      #first get it from the command line - there are some special cases to consider here 
      #since not all sh programs treat their host the same way
      if self.finargs != []:
        if self.section == "ssh" or self.section == "sftp":
          #strip out @ if it exists
          if "@" in self.finargs[0]:
            self.host = self.finargs[0].partition("@")[2]
          else:
            self.host = self.finargs[0]
          if self.section == "sftp" and ":" in self.host:
            self.host = self.host.partition(":")[0]
        elif self.section == "scp":
          #scp needs to be stripped similar to sftp, but for finargs[-1]
          #note this is limiting - if fwknop needs to be used on both hosts the
          #user will just have to do it by hand... but it seems in most cases the user will 
          #just want to copy something to a remote host???
          if "@" in self.finargs[-1]:
            self.host = self.finargs[-1].partition("@")[2]
          else:
            self.host = self.finargs[-1]
          if ":" in self.host:
            self.host = self.host.partition(":")[0]
        else:
          #by dafault - I'm tempted to just set self.host to self.finargs[0], but it probably should
          #be checked 
          self.usage()
      else:
        try:
          self.host = self.config.get("ssh_fwknop", "host")
        except:
          self.host = ""
      #if that doesn't exist get it from the config file
    if self.host == "":
      self.usage()
    return self.host  

#wrapper class - currently could be a functions :/
class ssh_fwknop():
  def __init__(self, clineoptions, configfile=None ):
    #host is a special option in the config file that can have a default place for 
    self.host = None
    self.verbocity = False

    self.sshcall = ""
    self.fwcall = ""

    #TODO: platform-independence
    #guess at possible configfile locations
    if configfile == None:
      if os.path.exists("./ssh_fwknop.conf"):
        configfile = "./ssh_fwknop.conf"
      elif os.path.exists("/etc/ssh_fwknop.conf"):
        configfile = "/etc/ssh_fwknop.conf"
      elif os.path.exists("/etc/fwknop/ssh_fwknop.conf"):
        configfile = "/etc/fwknop/ssh_fwknop.conf"
      else:
        if self.verbocity:
          sys.stdout.write("Configuration file not found. Continuing...\n")

    #default programs of ssh and fwknop
    sh_sect = "ssh"
    fw_sect = "fwknop"
    #search for alternate programs 
    #valid options are --scp, --sftp, --sshaskpass
    if "--scp" in clineoptions:
      clineoptions.remove("--scp")
      sh_sect = "scp"
    elif "--sftp" in clineoptions:
      clineoptions.remove("--sftp")
      sh_sect = "sftp"


    config = None
    if configfile != None:
      config = ConfigParser.ConfigParser()
      #override optionxform so config is case sensative
      config.optionxform = str
      config.read(configfile)
      #read the fwknop and shprog if they are defined
      try:
        shprog = config.get("ssh_fwknop", "sh_path")
        fwprog = config.get("ssh_fwknop", "fw_path")
      except:
        #TODO make system independent file shizaz and include the other possibilities
        #TODO also check that these exist and raise an error if they don't
        shprog = "/usr/bin/" + sh_sect
        fwprog = "/usr/bin/" + fw_sect
  
    sh_allopt = ssh_fwopts.alloptions(sh_sect)
    fw_allopt = ssh_fwopts.alloptions(fw_sect)
  
    #shprog options are called first so fwprog options can be stripped away
    shOpt =  getoptions(commandlineoptions=clineoptions, shortopt_args=sh_allopt.shortopts_arg, shortopt_noargs=sh_allopt.shortopts_noarg, 
             longopt_args=sh_allopt.longopts_arg, longopt_noargs=sh_allopt.longopts_noarg, binpath=shprog, config=config, section=sh_sect)
    self.sshcall =  shOpt.listify()
   
    #if verbocity is anywhere, we'll print extra information
    if "-v" in clineoptions:
      self.verbocity = True
    elif config != None:
      #get verbocity from the config file
      try:
        verbocity = config.get("ssh_fwknop", "verbose")
        if verbocity.lower() == "true":
          self.verbocity = True
      except:
        pass

    #get fwknop options
    fwclientopts = shOpt.get_fwknop()

    fwOpt = getoptions(commandlineoptions=fwclientopts, shortopt_args=fw_allopt.shortopts_arg, shortopt_noargs=fw_allopt.shortopts_noarg,
                       longopt_noargs=fw_allopt.longopts_noarg, longopt_args=fw_allopt.longopts_arg, binpath=fwprog, section=fw_sect, config=config, host=shOpt.gethost())
    self.fwcall = fwOpt.listify()
    #get fwknop host, but the -D rules all
    if "-D" not in self.fwcall:
      self.fwcall.append("-D")
      self.fwcall.append(shOpt.gethost())

    if self.verbocity:
      sys.stdout.write("Calling fw:\n " + " ".join(self.fwcall) + "\n")
    retcode = subprocess.call(self.fwcall)
    if retcode != 0:
      sys.stderr.write("FW ERROR\n")
      sys.exit(retcode)
    if self.verbocity:
      sys.stdout.write("Calling sh:\n" + " ".join(self.sshcall) + "\n")
    retcode = subprocess.call(self.sshcall)
    if retcode != 0:
      sys.stderr.write( "SSH ERROR\n")
      sys.exit(retcode)
 
if __name__ == "__main__":
  m = ssh_fwknop(sys.argv[1:])

